import numpy as np
import pandas as pd


class QlearningTable:
    def __init__(self, actions, learning_rate=0.01,reward_decay = 0.9,epsilon = 0.9):
        self.actions = actions
        self.q_table = pd.DataFrame(columns=self.actions, dtype=np.float64)
        self.epsilon = epsilon
        self.gamma = reward_decay
        self.lr = learning_rate


    def choose_action(self, observation):
        # 检查state状态在qtable中有没有，没有添加一下
        self.check_state_existe(observation)
        #  根据qtable 开始行为选择
        #  我们会让agent智能体选择最优的行为，所谓的最优其实actin在state下对应Q智最大
        #  如果我们一位的选择Q值最大的，就太贪婪了：你觉得宫保鸡丁好吃，你就一直吃宫保鸡丁，就不对其他的菜镜像尝试
        #  这是不太好的，因为少了很多尝试！万一比宫保鸡丁让你觉得更好吃的呢？比如佛跳墙
        #  所有我们需要再行为选择的时候多一下随机性，
        if np.random.uniform() < self.epsilon:
            state_actions = self.q_table.loc[observation, :]
            #  选择最优, 当行为的Q值数值一样大的时候随机选一个
            #  如果我们想要的是action的索引号， 而不是action对应的Q值，更关系的其实是最大Q值对应action是哪一个
            action = np.random.choice(state_actions[state_actions == np.max(state_actions)].index)
        else:
            # 最近选择一个行为， 其实就是进行了更多的尝试
            action = np.random.choice(self.actions)
        return action


    def learn(self, s, a, r, s_):
        # s, a, r, s_ 分别对应的就是 当前是的state状态，当前是个是action行为，环境接收到当前时刻action之后给出来的reward奖励， 下一时刻的state
        self.check_state_existe(s_) # 当前时刻的状态s已经在当前时刻选择行为的时候， 也就是当前时刻调用choose_action的时候已经检查过了
        if s is not None and a is not None:
            # 所谓的预测就是查询qtable里面的数值
            q_predict = self.q_table.loc[s, a]
            # 需要镜像判断，看看每一个回合是否进行了到了最后
            if s_ != 'terminal':
                # self.q_table.loc[s_, :] 代表 把下一刻时刻的所有action行为的Q值取出来
                q_target = r + self.gamma * self.q_table.loc[s_, :].max()
            else:
                q_target = r # 如果下一状态等于terminal
            # 更新
            self.q_table.loc[s, a] += self.lr * (q_target - q_predict)




    def check_state_existe(self,state):
        # 如果状态在qtable中不存在， 就添加一行
        if state not in self.q_table.index:
            new_row = pd.DataFrame(
                [[0] * len(self.actions)],
                columns=self.q_table.columns,
                index=[state]
            )
            self.q_table = pd.concat([self.q_table, new_row])
